
#include <thread>  // 线程库
#include <chrono>  // 时间库
#include <iostream>  // 标准输入输出流库
#include <cstdio>    // 标准输入输出库
#include <direct.h>   // 目录操作函数库
#include <vector>     // 动态数组（向量）库
#include "boost/filesystem/operations.hpp"  // Boost 文件系统库，用于文件和目录操作
#include "boost/filesystem/path.hpp"        // Boost 文件系统库，用于路径操作
#include "boost/filesystem/operations.hpp"  // 重复包含，可能是为了强调
#include "boost/filesystem/fstream.hpp"      // Boost 文件系统库，用于文件流操作
//#include "../txr_algo_dlm_v5/common/define.h"  // 包含通用定义
//#include "../txr_algo_dlm_v5/common/gpu_utility.h"  // 包含GPU工具函数
#include "common/device_utility.h"

#include "o_encode_dlm_v5.h"  // 包含模型转换函数
#include "model_yolo_v6.h"  // 包含YOLO v6类定义

/**
 * 枚举指定路径下所有匹配特定扩展名的文件。
 *
 * @param path 要枚举的目录路径。
 * @param extension 文件扩展名。
 * @return 返回一个包含所有匹配文件路径的向量。
 */
std::vector<std::string> EnumFiles(
    const std::string& path,
    const std::string &extension)
{
    namespace fs = boost::filesystem;  // 简化命名空间引用
    fs::path _path(path);  // 构造路径对象
    fs::path fullpath = fs::system_complete(_path);  // 获取系统的完整路径

    std::vector<std::string> ret;  // 用于存储结果的向量

    // 检查路径是否存在且为目录
    if (!fs::exists(fullpath) && !fs::is_directory(fullpath)) {
        return ret;  // 如果不是目录，返回空向量
    }
    fs::directory_iterator end_iter;  // 定义目录迭代器的结束迭代器
    for (fs::directory_iterator iter(fullpath); iter != end_iter; iter++) {  // 遍历目录
        try {
            if (fs::is_directory(*iter)) {  // 如果是目录，则跳过
                // ret.push_back(iter->path().wstring());  // 这行代码被注释掉了
            }
            // 如果是文件且扩展名匹配，则添加到结果向量
            else if (iter->path().extension() == extension) {
                ret.push_back(iter->path().string());
            }
        } catch (const std::exception&) {  // 捕获异常并跳过
            continue;
        }
    }
    return ret;  // 返回包含文件路径的向量
}

/**
 * 枚举模型配置文件并进行模型转换。
 *
 * 此函数遍历指定目录下的所有模型配置文件（.yaml），对于每个配置文件，
 * 它尝试找到对应的模型文件（.onnx或.ennx），然后使用YOLO v6类进行模型转换。
 * 转换成功后，将转换后的模型文件重命名并移动到目标文件夹。
 *
 * @param mon_folder 模型文件夹路径。
 * @param dest_folder 目标文件夹路径。
 * @param mark_info 标记信息，用于生成目标文件名。
 */
void enum_models(std::string& mon_folder, std::string& dest_folder, std::string & mark_info) {
    namespace fs = boost::filesystem;  // 简化命名空间引用
    std::string gpu_model = GetGPUModel();  // 获取GPU型号
    std::cout <<"gpu_model : " << gpu_model << std::endl;
    if (gpu_model.empty()) {  // 如果GPU型号未知
        std::cout << "Unknown Gpu Model! \n";  // 输出错误信息
        return;  // 返回
    }
    // 检查目标文件夹是否存在
    fs::path _dest_folder(dest_folder);
    if (!fs::exists(_dest_folder)) {  // 如果目标文件夹不存在，则返回
        return;
    }
    std::cout << _dest_folder.string() <<" exist"<< std::endl;
    std::vector<std::string> cfg_files = EnumFiles(mon_folder, std::string(".yaml"));  // 获取所有.yaml配置文件
    std::cout << "cfg_files : " << cfg_files.size() << std::endl;
    // 遍历所有配置文件
    for (const auto & cfg_file : cfg_files) {

        std::cout << "This cfg : " << cfg_file.c_str() << std::endl;
        st_trans_model_input input;  // 定义转换模型输入结构体
        fs::path _cfg_path(cfg_file);  // 获取配置文件的完整路径
        input.cfg_path = _cfg_path.string();  // 设置配置文件路径
        std::cout << "input.cfg_path : " << input.cfg_path << std::endl;

        fs::path _model_path;  // 定义模型文件路径变量
        // 尝试获取.ennx或.onnx模型文件路径
        {
            fs::path _path = _cfg_path.replace_extension(".ennx");
            if (fs::exists(_path)) {  // 如果.ennx文件存在
                _model_path = _path;  // 设置模型文件路径
            }
        }
        {
            fs::path _path = _cfg_path.replace_extension(".onnx");
            if (fs::exists(_path)) {  // 如果.onnx文件存在
                _model_path = _path;  // 设置模型文件路径
                std::cout << "model onnx: " << _model_path.string() << std::endl;
            }
        }
        if (_model_path.empty()) {  // 如果模型文件路径为空，则跳过
            continue;
        }

        fs::path _gdats_path(dest_folder + "\\" + _cfg_path.stem().string() + mark_info + ".hdats");  // 构造目标文件路径
        std::cout << "gdats_path : " << _gdats_path.string() << std::endl;
        if (fs::exists(_gdats_path)) {  // 如果目标文件已存在，则跳过
            continue;
        }

        // 输出开始信息
        std::cout << "START===============================================\n";
        // 设置模型文件路径
        input.model_path = _model_path.string();
        // 设置GPU型号
        input.gpu_model = gpu_model;
        // 设置转换文件路径（临时文件）
        input.trans_path = _gdats_path.string() + "-tmp";
        // 创建YOLO v6模型转换器实例
        yolo_v6 modelConverter;
        // 调用模型转换函数
        if (modelConverter.TransModel(input)) {
            // 输出完成信息
            std::cout << "DONE================================================\n";
            // 尝试重命名临时文件为最终文件名
            int renameResult = std::rename(
                input.trans_path.c_str(),
                _gdats_path.string().c_str());
            if (renameResult == 0) {
                // 如果重命名成功
                std::cout << "Conversion successful and file has been renamed.\n";
            } else {
                // 如果重命名失败
                std::cout << "Conversion successful, but file renaming failed.\n";
            }
        } else {
            // 如果转换失败
            std::cout << "FAIL!================================================\n";
        }
    }
}
int main(int argc, char* argv[]) {
    try {  // 尝试代码块
        char cur_folder[260];  // 定义字符数组，用于存储当前工作目录
        _getcwd(cur_folder, sizeof(cur_folder));  // 获取当前工作目录
        printf("Current working directory: \n %s\n", cur_folder);  // 输出当前工作目录

        std::string cwd = cur_folder;  // 将当前工作目录转换为字符串

        // 初始化模型文件夹和目标文件夹的路径为当前工作目录
        std::string mon_folder = cwd+"/model";
        std::string dest_folder = cwd;
        std::string mark_info = "v1";  // 默认标记信息

        // 根据命令行参数数量设置路径
        if (argc >= 2) {
            mon_folder = argv[1];  // 设置模型文件夹路径
        }
        if (argc >= 3) {
            dest_folder = argv[2];  // 设置目标文件夹路径
        }
        if (argc >= 4) {
            mark_info = argv[3];  // 设置标记信息
        }

        // 调用枚举模型并转换函数
        enum_models(mon_folder, dest_folder, mark_info);

        printf("Done!\n");  // 输出完成信息
    } catch (...) {  // 捕获所有异常
        printf("Exception!\n");  // 输出异常信息
    }

    return 0;  // 返回0表示程序正常退出
}
